# -*- coding: utf-8 -*-
"""
Created on Mon Nov 18 13:00:38 2013

@author: wudong

File location :
'D:/MyCodeLib/MyAbaqusLib/'

usage :
import drmHelper as dh

Reload the updated module :
reload
import drmHelper as dh
reload(ah)

"""
import numpy as np
import abaHelper as ah
import mdbHelper as mh
import odbHelper as oh
#
# ============================================================================
#
def outputModelMassStiffnessMatrix(model, elemSetName, outputFileName):
	instance = model.rootAssembly.instances.values()[0]
	instanceName = instance.name
	model.keywordBlock.synchVersions()
	mkDataOuputStr = """**
	** output mass and stiffness matrix
	**
	*element matrix output, elset=%s.%s, mass=yes, stiffness=yes, frequency=1, output file=user defined, file name=%s
	**""" % (instanceName, elemSetName, outputFileName)
	blockPrefix = '*End Step'
	insertPos = ah.getModelBlockPosition(model, blockPrefix)-1
	model.keywordBlock.insert(insertPos, mkDataOuputStr)
#
# ============================================================================
#
def outputModelInfNodeDispAcc(model, innerInfNodeSetName, outerInfNodeSetName, numIntervals):
	instance = model.rootAssembly.instances.values()[0]
	regionDef=instance.sets[innerInfNodeSetName]
	histName = 'Hist_%s' % (innerInfNodeSetName,)
	model.HistoryOutputRequest(name=histName, 
		createStepName=stepName, variables=('U1', 'U2', 'A1', 'A2'), region=regionDef, 
		numIntervals=numIntervals, sectionPoints=DEFAULT, rebar=EXCLUDE)
	#
	regionDef=instance.sets[outerInfNodeSetName]
	histName = 'Hist_%s' % (outerInfNodeSetName,)
	model.HistoryOutputRequest(name=histName, 
		createStepName=stepName, variables=('U1', 'U2', 'A1', 'A2'), region=regionDef, 
		numIntervals=numIntervals, sectionPoints=DEFAULT, rebar=EXCLUDE)
#
# ============================================================================
#
def getInfNodesAccDisp(odb, stepName, nodeSetName1, nodeSetName2, dofCount):
	nodeLabelsInner = oh.getOdbNodeLabelsFromNodeSet(odb, nodeSetName1)
	nodeLabelsOuter = oh.getOdbNodeLabelsFromNodeSet(odb, nodeSetName2)
	nodeLabels = nodeLabelsInner+nodeLabelsOuter
	nodeCount = len(nodeLabels)
	#
	accVarNames = ['A1', 'A2']
	dispVarNames = ['U1', 'U2']
	assert ((dofCount==2) or (dofCount==3))
	if dofCount==3:
		accVarNames = accVarNames+['A3']
		dispVarNames = dispVarNames+['U3']
	totalDofCount = nodeCount*dofCount
	#
	accDataList = []
	dispDataList = []
	for nodeLabel in nodeLabels:
		for varName in accVarNames:
			data = oh.getOdbNodeHistoryOutputData(odb, stepName, nodeLabel, varName)
			accDataList.append(data)
		#
		for varName in dispVarNames:
			data = oh.getOdbNodeHistoryOutputData(odb, stepName, nodeLabel, varName)
			dispDataList.append(data)
	assert ((len(accDataList)==totalDofCount) and (len(dispDataList)==totalDofCount))
	#
	timePointData = accDataList[0][:, 0]
	timePointCount = accDataList[0].shape[0]
	assert (len(accDataList)==len(dispDataList))
	for i in range(totalDofCount):
		assert (accDataList[i].shape[0]==timePointCount) and (dispDataList[i].shape[0]==timePointCount)
	#
	accData = np.zeros(shape=(totalDofCount, timePointCount))
	dispData = np.zeros(shape=(totalDofCount, timePointCount))
	for i in range(totalDofCount):
		accData[i, :] = np.reshape(accDataList[i][:, 1], (1,timePointCount))
		dispData[i, :] = np.reshape(dispDataList[i][:, 1], (1,timePointCount))
	#
	return (nodeLabelsInner, nodeLabelsOuter, timePointData, accData, dispData)
#
# ============================================================================
#
def parseIntNumberLines(lines):
	nums = []
	for line in '  '.join(lines).replace('*', ' ').replace(',', ' ').split():
		l = line.strip()
		if (len(l)>0):
			nums.append(int(l))
	#
	return nums
#
# ============================================================================
#
def parseFloatNumberLines(lines):
	nums = []
	for line in '  '.join(lines).replace('*', ' ').replace(',', ' ').split():
		l = line.strip()
		if (len(l)>0):
			nums.append(float(l))
	#
	return nums
#
# ============================================================================
#
def parseMatrixLines(lines, nodeCount, dofCount):
	nums = parseFloatNumberLines(lines)
	#
	totalDofCount = nodeCount*dofCount
	assert (len(nums)==(totalDofCount*(totalDofCount+1)/2))
	#
	mkshape=(totalDofCount,totalDofCount)
	mm = np.zeros(mkshape)
	istart = 0
	for n in range(1,mkshape[0]+1):
		mm[n-1,0:n] = np.array(nums[istart:(istart+n)])
		istart = istart+n
	#
	mm2 = mm
	for i in range(0,mkshape[0]):
		mm2[i,i] = mm2[i,i]*0.5
	#
	mm = mm2+np.transpose(mm2)
	return mm
#
# ============================================================================
#
def findFirstLineIndex(lines, startString, startIndex=0):
	firstPos = 0
	for i in range(startIndex, len(lines)):
		if lines[i].startswith(startString):
			firstPos = i
			break
	#
	return firstPos
#
# ============================================================================
#
def parseMkDataFile(mkDataFileName):
	f = open(mkDataFileName, 'r')
	allLines = f.readlines()
	f.close()
	#
	firstPos = findFirstLineIndex(allLines, '** ELEMENT NUMBER')
	secondPos = findFirstLineIndex(allLines, '** ELEMENT NUMBER', startIndex=(firstPos+1))
	#
	lineCountPerGroup = secondPos-firstPos
	groupCount = int(len(allLines)/lineCountPerGroup)
	#
	firstGroupLines = allLines[0:lineCountPerGroup]
	elemNumLineIndex = findFirstLineIndex(firstGroupLines, '** ELEMENT NUMBER')
	nodeCountLineIndex = findFirstLineIndex(firstGroupLines, '*USER ELEMENT')
	nodeNumLineIndex = findFirstLineIndex(firstGroupLines, '** ELEMENT NODES')
	stiffMatrixLineIndex = findFirstLineIndex(firstGroupLines, '*MATRIX,TYPE=STIFFNESS')
	massMatrixLineIndex = findFirstLineIndex(firstGroupLines, '*MATRIX,TYPE=MASS')
	dofLineIndex = stiffMatrixLineIndex-1
	#
	elemNums = []
	for i in range(groupCount):
		groupStart = lineCountPerGroup*i
		istart = groupStart+elemNumLineIndex
		line = allLines[istart]
		elemNum = int(line.split()[3])
		elemNums.append(elemNum)
	#
	incCount = elemNums.count(elemNums[0])		# increment count
	elemCount = int(groupCount/incCount)		# real elem count
	assert ((elemCount*incCount*lineCountPerGroup)==len(allLines))
	#
	ens = elemNums[0:elemCount] 	# element number list
	nns = []    # node number list
	mms = []    # mass matrix list
	kms = []    # stiffness matrix list
	for i in range(elemCount):      # for each element
		groupStart = lineCountPerGroup*i
		#
		# elem number
		#
		istart = groupStart+elemNumLineIndex
		line = allLines[istart]
		en = int(line.split()[3])
		assert (en==ens[i])
		#
		# node count for current elem
		#
		istart = groupStart+nodeCountLineIndex
		line = allLines[istart]
		nodeCount = int(line.replace(',', ' ').split()[3])
		#
		# node number list
		#
		istart = groupStart+nodeNumLineIndex+1
		iend = groupStart+dofLineIndex
		lines = [line.strip() for line in allLines[istart:iend]]
		nodeNums = parseIntNumberLines(lines)
		assert (len(nodeNums)==nodeCount)
		nns.append(nodeNums)
		#
		# dof number list
		#
		istart = groupStart+dofLineIndex
		iend = groupStart+stiffMatrixLineIndex
		lines = [line.strip() for line in allLines[istart:iend]]
		dofNums = parseIntNumberLines(lines)
		dofCount = len(dofNums)
		#
		# stiffness matrix
		#
		istart = groupStart+stiffMatrixLineIndex+1
		iend = groupStart+massMatrixLineIndex
		lines = [line.strip() for line in allLines[istart:iend]]
		km = parseMatrixLines(lines, nodeCount, dofCount)
		kms.append(km)
		#
		# mass matrix
		#
		istart = groupStart+massMatrixLineIndex+1
		iend = groupStart+lineCountPerGroup
		lines = [line.strip() for line in allLines[istart:iend]]
		mm = parseMatrixLines(lines, nodeCount, dofCount)
		mms.append(mm)
	#
	return (ens, nns, mms, kms, dofCount)
#
# ============================================================================
#
def createMKbeMatrix(ens, nns, mms, kms, nodeLabels, dofCount):
	totalNodeCount = len(nodeLabels)	# nodeLabels = nodeLabelsInner + nodeLabelsOuter
	totalDofCount = totalNodeCount*dofCount
	mkshape=(totalDofCount,totalDofCount)
	Mmbe = np.zeros(shape=mkshape)
	Kmbe = np.zeros(shape=mkshape)
	elemCount = len(ens)
	for k in range(0,elemCount):	# for each element
		elemNum = ens[k]			# element number of current element
		nodeNums = nns[k]			# node number list of current element
		mapIndex = []				# position of each node for current element in nodeLabels
		#for i in range(0,len(nodeNums)):
		#	for j in range(0,len(nodeLabels)):
		#		if nodeNums[i]==nodeLabels[j]:
		#			mapIndex.append(j)
		for nodeNum in nodeNums:
			assert nodeNum in nodeLabels, 'nodeNum=%d, elemNum=%d' % (nodeNum, elemNum)
			mapIndex.append(nodeLabels.index(nodeNum))
		#
		assert len(mapIndex)==len(nodeNums), ('len(mapIndex)=%d, len(nodeNums)=%d, nodeNums[0]=%d'%(len(mapIndex), len(nodeNums), nodeNums[0]))
		mm = mms[k]
		km = kms[k]
		for i in range(0,len(nodeNums)):
			for j in range(0,len(nodeNums)):
				row = mapIndex[i]
				col = mapIndex[j]
				#
				rowStart = row*dofCount
				rowEnd = rowStart+dofCount
				colStart = col*dofCount
				colEnd = colStart+dofCount
				iStart = i*dofCount
				iEnd = iStart+dofCount
				jStart = j*dofCount
				jEnd = jStart+dofCount
				#
				Mmbe[rowStart:rowEnd, colStart:colEnd] = Mmbe[rowStart:rowEnd, colStart:colEnd]+mm[iStart:iEnd, jStart:jEnd]
				Kmbe[rowStart:rowEnd, colStart:colEnd] = Kmbe[rowStart:rowEnd, colStart:colEnd]+km[iStart:iEnd, jStart:jEnd]
	#
	return (Mmbe, Kmbe)
#
# ============================================================================
#
def calcPeb(Mmbe, Kmbe, accData, dispData, nodeLabelsInner, nodeLabelsOuter, dofCount):
	timePointCount = accData.shape[1]
	assert timePointCount==dispData.shape[1]
	nodeCountInner = len(nodeLabelsInner)
	nodeCount = len(nodeLabelsInner)+len(nodeLabelsOuter)
	totalDofCount = nodeCount*dofCount
	assert (accData.shape[0]==totalDofCount) and (dispData.shape[0]==totalDofCount)
	Peb = np.zeros(shape=(totalDofCount, timePointCount))
	beIndex = nodeCountInner*dofCount
	#
	Mbe = Mmbe[0:beIndex, beIndex:]
	ae0 = accData[beIndex:, :]
	Kbe = Kmbe[0:beIndex, beIndex:]
	ue0 = dispData[beIndex:, :]
	Pb = -np.dot(Mbe,ae0)-np.dot(Kbe,ue0)
	#
	Meb = Mmbe[beIndex:, 0:beIndex]
	ab0 = accData[0:beIndex, :]
	Keb = Kmbe[beIndex:, 0:beIndex]
	ub0 = dispData[0:beIndex, :]
	Pe = np.dot(Meb,ab0)+np.dot(Keb,ub0)
	#
	assert (Pb.shape[0]==beIndex) and (Pe.shape[0]==(totalDofCount-beIndex))
	Peb[0:beIndex, :] = Pb
	Peb[beIndex:, :] = Pe
	#
	return Peb
#
# ============================================================================
#
def calcPebFromAccDispData(mkDataFileName, accData, dispData, nodeLabelsInner, nodeLabelsOuter):
	assert accData.shape==dispData.shape
	nodeLabels = nodeLabelsInner+nodeLabelsOuter
	(ens, nns, mms, kms, dofCount) = parseMkDataFile(mkDataFileName)
	(Mmbe, Kmbe) = createMKbeMatrix(ens, nns, mms, kms, nodeLabels, dofCount)
	Peb = calcPeb(Mmbe, Kmbe, accData, dispData, nodeLabelsInner, nodeLabelsOuter, dofCount)
	return Peb
#
# ============================================================================
#
def applyModelPeb(model, createStepName, nodeLabels, timePointData, Peb):
	nodeCount = len(nodeLabels)
	totalDofCount = Peb.shape[0]
	assert totalDofCount%nodeCount==0
	dofCount = totalDofCount/nodeCount
	timePointCount = timePointData.shape[0]
	assert timePointCount==Peb.shape[1]
	#
	ampData = np.zeros(shape=(timePointCount, 2))
	ampData[:,0] = timePointData.reshape((timePointCount,))
	for i in range(nodeCount):
		nodeLabel = int(nodeLabels[i])
		for j in range(dofCount):
			ampData[:,1] = Peb[i*dofCount+j,:].reshape((timePointCount,))
			dof = j+1
			mh.applyModelDynamicConcentratedForce(model, createStepName, nodeLabel, dof, ampData)
#
# ============================================================================
#
# Get node labels of DRM interface nodes
#
def getInterfaceNodeLabels(odb, drmInfInnerNodeSetName, drmInfOuterNodeSetName):
	instance = odb.rootAssembly.instances.values()[0]
	nodeSetName = drmInfInnerNodeSetName	# b
	nodeLabelsInner = []
	for node in instance.nodeSets[nodeSetName].nodes:
		nodeLabelsInner.append(node.label)
	#
	nodeSetName = drmInfOuterNodeSetName	# e
	nodeLabelsOuter = []
	for node in instance.nodeSets[nodeSetName].nodes:
		nodeLabelsOuter.append(node.label)
	#
	nodeLabels = nodeLabelsInner+nodeLabelsOuter	# nodeLabels = [nodeLabelsInner, nodeLabelsOuter]
	return (nodeLabels, nodeLabelsInner, nodeLabelsOuter)
#
# ============================================================================
# Get acceleration and displacement time history of DRM interface nodes from the global model
#
def getInterfaceNodeAccDispData(odb, stepName, nodeLabels):
	dispVarNames = ['U1', 'U2']
	step = odb.steps[stepName]
	regionName = 'Node PART_ANSMODEL-1.%d' % (nodeLabels[0])
	regionDef = step.historyRegions[regionName]
	data = regionDef.historyOutputs[dispVarNames[0]].data
	timePointCount = len(data)
	data = ah.tableToArray(data)
	timePointData = data[:,0]
	#
	nodeCount = len(nodeLabels)
	varCount = len(dispVarNames)
	dispData = np.zeros(shape=(nodeCount*varCount, timePointCount))
	for i in range(nodeCount):
		nodeLabel = nodeLabels[i]
		for j in range(varCount):
			varName = dispVarNames[j]
			regionName = 'Node PART_ANSMODEL-1.%d' % (nodeLabel)
			regionDef = step.historyRegions[regionName]
			data = ah.tableToArray(regionDef.historyOutputs[varName].data)
			dispData[i*varCount+j, :] = np.reshape(data[:, 1], (1,timePointCount))
	#
	accVarNames = ['A1', 'A2']
	assert varCount == len(accVarNames)
	accData = np.zeros(shape=(nodeCount*varCount, timePointCount))
	for i in range(nodeCount):
		nodeLabel = nodeLabels[i]
		for j in range(varCount):
			varName = accVarNames[j]
			regionName = 'Node PART_ANSMODEL-1.%d' % (nodeLabel)
			regionDef = step.historyRegions[regionName]
			data = ah.tableToArray(regionDef.historyOutputs[varName].data)
			accData[i*varCount+j, :] = np.reshape(data[:, 1], (1,timePointCount))
	#
	return (timePointData,accData,dispData)
#
# ============================================================================
#
# Ensure the number and coordinate of interface nodes in local model coincide with that in global model
#
def calcPebFromOdbFile(odbFilePath, stepName, drmInfInnerNodeSetName, drmInfOuterNodeSetName, mkDataFilePath):
	odb = oh.openOdb(odbFilePath, readOnly=True)
	(nodeLabels, nodeLabelsInner, nodeLabelsOuter) = getInterfaceNodeLabels(odb, drmInfInnerNodeSetName, drmInfOuterNodeSetName)
	(timePointData,accData,dispData) = getInterfaceNodeAccDispData(odb, stepName, nodeLabels)
	odb.close()
	#
	(ens, nns, mms, kms, dofCount) = parseMkDataFile(mkDataFilePath)
	(Mmbe, Kmbe) = createMKbeMatrix(ens, nns, mms, kms, nodeLabels, dofCount)
	Peb = calcPeb(Mmbe, Kmbe, accData, dispData, nodeLabelsInner, nodeLabelsOuter, dofCount)
	return (nodeLabels, timePointData, Peb)
#
# ============================================================================
#
# Ensure the y coordinate of interface nodes in local model coincide with that in global model
#
def calcPebFromSoilColumnOdbFile(odbFilePath, stepName, resultNodeSetName, abaMdb, drmInfInnerNodeSetName, drmInfOuterNodeSetName, mkDataFileName):
	accVarNames = ['A1', 'A2']
	dispVarNames = ['U1', 'U2']
	dofCount = len(accVarNames)
	accDataList = []
	dispDataList = []
	yCoords = []
	abaOdb = oh.AbaqusOdb(odbFilePath)
	abaOdb.openOdb()
	nodeLabels = abaOdb.getNodeLabelsFromNodeSet(resultNodeSetName)
	for nodeLabel in nodeLabels:
		yCoord = abaOdb.getNodeCoordinateFromNodeLabel(nodeLabel)[1]
		yCoords.append(yCoord)
		for varName in accVarNames:
			accData = abaOdb.getNodeHistoryOutput(stepName, nodeLabel, varName)
			accDataList.append(accData)
	
		for varName in dispVarNames:
			dispData = abaOdb.getNodeHistoryOutput(stepName, nodeLabel, varName)
			dispDataList.append(dispData)
	
	abaOdb.closeOdb()
	assert len(accDataList)==(len(nodeLabels)*dofCount) and len(dispDataList)==(len(nodeLabels)*dofCount)
	#
	nodeLabelsInner = abaMdb.getPartNodeLabelsFromNodeSet(drmInfInnerNodeSetName)
	nodeLabelsOuter = abaMdb.getPartNodeLabelsFromNodeSet(drmInfOuterNodeSetName)
	nodeLabels = nodeLabelsInner+nodeLabelsOuter
	#
	timePointData = accDataList[0][:,0]
	timePointCount = timePointData.shape[0]
	for i in range(len(accDataList)):
		assert timePointCount==accDataList[i].shape[0] and timePointCount==dispDataList[i].shape[0], 'i = %d' % (i,)

	nodeCount = len(nodeLabels)
	accData = np.zeros(shape=(nodeCount*dofCount, timePointCount))
	dispData = np.zeros(shape=(nodeCount*dofCount, timePointCount))
	assert len(accDataList)==len(dispDataList)
	eps = 0.1
	for i in range(nodeCount):				# for each node of interface element in free field model
		nodeLabel = nodeLabels[i]
		nodeCoord = abaMdb.getPartNodeCoordinateFromNodeLabel(nodeLabel)
		mapIndex = -1
		for j in range(len(yCoords)):		# for each y coordinate of node in soil column model
			if abs(nodeCoord[1]-yCoords[j])<eps:
				mapIndex = j
				break
		#
		assert mapIndex>=0
		for j in range(dofCount):
			accData[i*dofCount+j, :] = np.reshape(accDataList[mapIndex*dofCount+j][:, 1], (1, timePointCount))
			dispData[i*dofCount+j, :] = np.reshape(dispDataList[mapIndex*dofCount+j][:, 1], (1, timePointCount))
		
		#j = 0
		#accData[i*dofCount+j, :] = np.reshape(accDataList[mapIndex*dofCount+j][:, 1], (1, timePointCount))
		#dispData[i*dofCount+j, :] = np.reshape(dispDataList[mapIndex*dofCount+j][:, 1], (1, timePointCount))
	
	#
	(ens, nns, mms, kms, dofCount) = parseMkDataFile(mkDataFileName)
	assert dofCount==len(accVarNames)
	for nodeNums in nns:
		for nodeNum in nodeNums:
			assert nodeNum in nodeLabels
	
	#
	(Mmbe, Kmbe) = createMKbeMatrix(ens, nns, mms, kms, nodeLabels, dofCount)
	Peb = calcPeb(Mmbe, Kmbe, accData, dispData, nodeLabelsInner, nodeLabelsOuter, dofCount)
	return (nodeLabels, timePointData, Peb)
#
# ============================================================================
#